import os
import shutil
from tqdm import tqdm


def safe_mkdir(d):
    if not os.path.exists(d):
        os.mkdir(d)


def group_avi_by_channel(root_src, chn, root_dst):
    '''
        将同属一个通道的avi视频片段放入同名文件夹（如必要则创建）
    '''
    fp_list = []
    if not os.path.isdir(root_dst):
        safe_mkdir(root_dst)
    for f in os.listdir(root_src):
        if f.endswith('.avi') and int(f.split('-')[0]) == chn:
            fp_list.append(os.path.join(root_src, f))
    [shutil.move(fp, root_dst) for fp in tqdm(fp_list, ascii=True)]
        

def scan_ch_prefix(root_path):
    '''
        将文件夹下所有avi文件按通道前缀分类，并返回通道号列表
    '''
    chn_list = []
    for f in os.listdir(root_path):
        if f.endswith('.avi'):
            chn = int(f.split('-')[0])
            if chn not in chn_list:
                chn_list.append(chn)
    return chn_list


def is_contain_file(d, f_name):
    if not os.path.isdir(d):
        raise Exception('Invalid directory %s' % d)
    for f in os.listdir(d):
        if f == f_name:
            return True


def is_contain_file_like(d, f_suffix):
    if not os.path.isdir(d):
        raise Exception('Invalid directory %s' % d)
    for f in os.listdir(d):
        if f.endswith(f_suffix):
            return True


def list_avi_paths(root_path):
    '''
        列出所有包含avi视频的通道文件夹绝对路径
    '''
    ch_plist = []
    for d in os.listdir(root_path):
        dpath = os.path.join(root_path, d)
        if not os.path.isdir(dpath):
            continue
        if is_contain_file_like(dpath, '.avi'):
            ch_plist.append(dpath)
    return ch_plist


def load_classes(path):
    """
    Loads class labels at 'path'
    """
    names = []
    with open(path, "r") as fp:
        for line in fp.readlines():
            line = line.rstrip('\r\t\n ')
            if len(line) > 0:
                names.append(line)
    return names


if __name__ == '__main__':
    # 测试视频分组 
    root_src = '/home/youngway/data/mtrace/20201202_c57-males-dvr1'
    chn = 10
    root_dst = os.path.join(root_src, str(chn))
    group_avi_by_channel(root_src, chn, root_dst)
    # 测试通道列表
    chlist = list_avi_paths(root_src)
    for ch in chlist:
        if is_contain_file(ch, 'locations.txt'):
            print('processed ch path: %s' % ch)
    # print(chlist)